/**
 * RDDL: A simple graphics display for the Sidewalk domain. 
 * 
 * @author Scott Sanner (ssanner@gmail.com)
 * @version 10/10/10
 *
 **/

package br.usp.pcs.lti.planning;

import java.awt.Color;
import java.util.ArrayList;
import java.util.Map;

import rddl.EvalException;
import rddl.RDDL.LCONST;
import rddl.RDDL.PVAR_NAME;
import rddl.RDDL.TYPE_NAME;
import rddl.State;
import rddl.viz.BlockDisplay;
import rddl.viz.StateViz;

public class TrafficDisplay extends StateViz {

    public TrafficDisplay() {
        this._nTimeDelay = 0; // in milliseconds
    }

    public TrafficDisplay(int time_delay_per_frame) {
        this._nTimeDelay = time_delay_per_frame; // in milliseconds
    }

    public boolean _bSuppressNonFluents = true;
    public BlockDisplay _bd = null;
    public int _nTimeDelay = 0;

    public void display(State s, int time) {
        System.out.println("TIME = " + time + ": " + this.getStateDescription(s));
    }

    // ////////////////////////////////////////////////////////////////////

    public String getStateDescription(State s) {
        StringBuilder sb = new StringBuilder();

        this.animate(s);

        // Go through all variable types (state, interm, observ, action, nonfluent)
        for (Map.Entry<String, ArrayList<PVAR_NAME>> e : s._hmTypeMap.entrySet()) {

            if (this._bSuppressNonFluents && e.getKey().equals("nonfluent"))
                continue;
            if (e.getKey().equals("interm"))
                continue;

            // Go through all variable names p for a variable type
            for (PVAR_NAME p : e.getValue())
                try {
                    // Go through all term groundings for variable p
                    ArrayList<ArrayList<LCONST>> gfluents = s.generateAtoms(p);
                    for (ArrayList<LCONST> gfluent : gfluents) {
                        Object value = s.getPVariableAssign(p, gfluent);
                        if (e.getKey().equals("action") || value.equals(Boolean.TRUE)
                                || p._sPVarName.startsWith("light-signal")) {
                            sb.append("- " + e.getKey() + ": " + p + (gfluent.size() > 0 ? gfluent : "") + " := "
                                    + value + "\n");
                        }
                    }

                } catch (EvalException ex) {
                    sb.append("- could not retrieve assignment" + s + " for " + p + "\n");
                }
        }

        // Sleep so the animation can be viewed at a frame rate of 1000/_nTimeDelay per second
        try {
            Thread.sleep(this._nTimeDelay);
        } catch (InterruptedException e) {
            System.err.println(e);
            e.printStackTrace(System.err);
        }

        return sb.toString();
    }

    private void animate(State s) {
        TYPE_NAME cell_type = new TYPE_NAME("cell");
        ArrayList<LCONST> cells = s._hmObject2Consts.get(cell_type);

        TYPE_NAME vehicle_type = new TYPE_NAME("vehicle");
        ArrayList<LCONST> vehicles = s._hmObject2Consts.get(vehicle_type);

        TYPE_NAME intersection_type = new TYPE_NAME("intersection");
        ArrayList<LCONST> intersections = s._hmObject2Consts.get(intersection_type);

        PVAR_NAME at = new PVAR_NAME("at");
        PVAR_NAME ambulance = new PVAR_NAME("AMBULANCE");
        PVAR_NAME bus = new PVAR_NAME("BUS");
        PVAR_NAME light_signal1 = new PVAR_NAME("light-signal1");
        PVAR_NAME light_signal2 = new PVAR_NAME("light-signal2");

        if (this._bd == null) {
            int max_row = -1;
            int max_col = -1;
            for (LCONST cell : cells) {
                String[] split = cell.toString().split("[_a]");
                int row = new Integer(split[1]);
                int col = new Integer(split[2]);
                if (row > max_row)
                    max_row = row;
                if (col > max_col)
                    max_col = col;
            }
            for (LCONST intersection : intersections) {
                String[] split = intersection.toString().split("[_a]");
                if (split.length != 3)
                    continue;

                int row = new Integer(split[1]);
                int col = new Integer(split[2]);
                if (row > max_row)
                    max_row = row;
                if (col > max_col)
                    max_col = col;
            }

            this._bd = new BlockDisplay("RDDL Traffic Simulation", "RDDL Traffic Simulation", max_row + 2, max_col + 2);
        }

        // Set up an arity-1 parameter list
        ArrayList<LCONST> params = new ArrayList<LCONST>(1);
        params.add(null);
        // Set up an arity-2 parameter list
        ArrayList<LCONST> params2 = new ArrayList<LCONST>(2);
        params2.add(null);
        params2.add(null);

        this._bd.clearAllCells();
        this._bd.clearAllLines();
        for (LCONST cell : cells) {

            // System.out.println("'" + cell + "'");
            String[] split = cell.toString().split("[_a]");
            int row = new Integer(split[1]);
            int col = new Integer(split[2]);

            Boolean occ = false;
            String symbol = null;
            Color c = Color.gray; // _bd._colors[row % _bd._colors.length];
            for (LCONST vehicle : vehicles) {
                // Get cell info and select a color
                params2.set(0, vehicle);
                params2.set(1, cell);
                boolean vehicleAtCell = (Boolean) s.getPVariableAssign(at, params2);

                if (vehicleAtCell) {
                    occ = true;
                    symbol = "" + vehicle.toString().charAt(vehicle.toString().length() - 1);

                    params.set(0, vehicle);
                    if ((Boolean) s.getPVariableAssign(ambulance, params)) {
                        c = Color.RED;
                    } else if ((Boolean) s.getPVariableAssign(bus, params)) {
                        c = Color.BLUE;
                    }
                    break;
                }

            }
            // System.out.println(occupied + " (" + row + "," + col + ") = " + occ);

            this._bd.setCell(row, col, c, occ ? symbol : ".");
        }

        for (LCONST intersection : intersections) {
            params.set(0, intersection);
            String[] split = intersection.toString().split("[_a]");
            if (split.length != 3)
                continue;

            int row = new Integer(split[1]);
            int col = new Integer(split[2]);
            boolean b_ls1 = (Boolean) s.getPVariableAssign(light_signal1, params);
            boolean b_ls2 = (Boolean) s.getPVariableAssign(light_signal2, params);
            boolean red = (b_ls1 && b_ls2) || (!b_ls1 && !b_ls2);
            if (red) {
                this._bd.setCell(row, col, Color.red, null);
                // _bd.setCell(row+1, col, Color.red, null);
            } else {
                boolean ns = (!b_ls1 && b_ls2);
                this._bd.setCell(row, col, Color.green, ns ? "||" : "--");
                // _bd.setCell(row+1, col, Color.green, ns ? "||" : "--");
            }

        }

        this._bd.repaint();
    }

    public void close() {
        if (this._bd != null) {
            this._bd.close();
        }
    }
}
